#!/usr/bin/env python
# coding: utf-8

"""
Python module to download the MNIST dataset and split the dataset into 5 tasks to bring the total to
9 tasks in all for the sequence

"""
import warnings
import os
from pathlib import Path
import shutil
import gdown
import pandas as pd
import requests, zipfile, io
import time

since = time.time()

# path to the data file, already been created
path_to_file = "../Data"
root_path = "../images/"
train_path_oracle = "../images/image1"
test_path_oracle = "../images/image2"
train_source_path = "../images/source.csv"
test_target_path = "../images/target.csv"

dir_list = next(os.walk(train_path_oracle))[1]
dir_list = [[dir_list[0], dir_list[1]], [dir_list[2], dir_list[3], dir_list[4]]]
# dictionary mapping the digits to the task directory
# images belonging to '0' and '1' map to Task_5 directory
# images belonging to '2' and '3' map to Task_6 directory
# and so on

task_dict = {
    dir_list[0][0]: '1',
    dir_list[0][1]: '1',
    dir_list[1][0]: '2',
    dir_list[1][1]: '2',
    dir_list[1][2]: '2'
             }

# create directories for all the tasks
try:
    for digits in dir_list:
        digit_0 = digits[0]
        digit_1 = digits[1]
        dir_key = task_dict[digit_0]

        # create the directory
        task_dir = os.path.join(path_to_file, "Task_" + dir_key)
        if not os.path.exists(task_dir):
            os.makedirs(task_dir)

        # create the train and test sub directories within the Task folder
        train_path = os.path.join(task_dir, "train")
        test_path = os.path.join(task_dir, "test")

        os.mkdir(train_path)
        os.mkdir(test_path)

        # create the class directories within the test and train folders
        train_dest_0 = os.path.join(train_path, digit_0)
        train_dest_1 = os.path.join(train_path, digit_1)
        test_dest_0 = os.path.join(test_path, digit_0)
        test_dest_1 = os.path.join(test_path, digit_1)
        os.mkdir(train_dest_0)
        os.mkdir(train_dest_1)
        os.mkdir(test_dest_0)
        os.mkdir(test_dest_1)

        if len(digits) == 3:
            digit_2 = digits[2]
            train_dest_2 = os.path.join(train_path, digit_2)
            test_dest_2 = os.path.join(test_path, digit_2)
            os.mkdir(train_dest_2)
            os.mkdir(test_dest_2)

except FileExistsError:
    pass

# move the files to the appropriate folders
for folder in next(os.walk(train_path_oracle))[1]:
    train_files_path = os.path.join(train_path_oracle, folder)
    test_files_path = os.path.join(test_path_oracle, folder)

    # 将训练文件复制至对应文件夹
    source_paths = pd.read_csv(train_source_path)
    source_paths = source_paths[source_paths['label'] == folder]['Path']
    for source_p in source_paths:
        path = os.path.join(root_path, source_p)
        if os.path.isfile(path):
            shutil.copy(path, os.path.join(path_to_file + "/Task_" + task_dict[folder] + f"/train/{folder}/", source_p.split("/")[1]))

    # 将测试文件复制至对应文件夹
    target_paths = pd.read_csv(test_target_path)
    target_paths = target_paths[target_paths['label'] == folder]['Path']
    for target_p in target_paths:
        path = os.path.join(root_path, target_p)
        if os.path.isfile(path):
            shutil.copy(path, os.path.join(path_to_file + "/Task_" + task_dict[folder] + f"/test/{folder}/", target_p.split("/")[1]))

    for filename in os.listdir(train_files_path):
        train_file_path = os.path.join(train_files_path, filename)
        if os.path.isfile(train_file_path):
            shutil.copy(train_file_path, os.path.join(path_to_file + "/Task_" + task_dict[folder] + f"/train/{folder}/", filename))

    for filename in os.listdir(test_files_path):
        test_file_path = os.path.join(test_files_path, filename)
        if os.path.isfile(test_file_path):
            shutil.copy(test_file_path, os.path.join(path_to_file + "/Task_" + task_dict[folder] + f"/test/{folder}/", filename))

total_time = time.time() - since

print("This process took {} minutes and {} seconds to execute".format(total_time // 60, total_time % 60))